import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris

class KMeans:
    def __init__(self, n_clusters=3, max_iter=100, random_state=42):
        self.n_clusters = n_clusters
        self.max_iter = max_iter
        self.random_state = random_state
        self.centers = None
        self.labels = None

    def _euclidean_distance(self, x1, x2):
        return np.sqrt(np.sum((x1 - x2)**2, axis=1))

    def fit(self, X):
        np.random.seed(self.random_state)
        n_samples, n_features = X.shape

        random_indices = np.random.choice(n_samples, size=self.n_clusters, replace=False)
        self.centers = X[random_indices]

        for _ in range(self.max_iter):
            distances = np.array([self._euclidean_distance(X, center) for center in self.centers]).T
            self.labels = np.argmin(distances, axis=1)

            new_centers = np.array([X[self.labels == i].mean(axis=0) for i in range(self.n_clusters)])

            if np.allclose(self.centers, new_centers, atol=1e-6):
                break

            self.centers = new_centers

        self.sse = np.sum([np.sum((X[self.labels == i] - self.centers[i])**2) for i in range(self.n_clusters)])
        return self.labels, self.centers, self.sse

def load_iris_dataset():
    iris = load_iris()
    X = iris.data[:, :2]
    y = iris.target
    feature_names = iris.feature_names[:2]
    target_names = iris.target_names
    return X, y, feature_names, target_names

def plot_clusters(X, labels, centers, feature_names):
    plt.figure(figsize=(8, 6))
    colors = ['red', 'blue', 'green']
    markers = ['o', 's', '^']
    for i in range(len(centers)):
        cluster_samples = X[labels == i]
        plt.scatter(cluster_samples[:, 0], cluster_samples[:, 1], c=colors[i], marker=markers[i], label=f'簇 {i+1}', alpha=0.6)
    plt.scatter(centers[:, 0], centers[:, 1], c='black', marker='x', s=200, linewidths=3, label='聚类中心')
    plt.xlabel(feature_names[0])
    plt.ylabel(feature_names[1])
    plt.title('k-means聚类结果')
    plt.legend()
    plt.savefig('kmeans_clusters.png')
    plt.show()

if __name__ == "__main__":
    X, y_true, feature_names, target_names = load_iris_dataset()
    kmeans = KMeans(n_clusters=3, max_iter=100, random_state=42)
    labels, centers, sse = kmeans.fit(X)
    plot_clusters(X, labels, centers, feature_names)
    print(f"k-means聚类完成，簇内平方和（SSE）：{sse:.4f}")
    print(f"最终聚类中心：\n{centers}")